import os
import cv2
import numpy as np
import random
from tqdm import tqdm
import re
# 在文件开头或合适位置，添加加载PCB背景图路径列表
pcb_bg_folder = "/home/lhx/pcbsyn/data_zoo/jrstruecolordepthbg"
pcb_bg_paths = [os.path.join(pcb_bg_folder, f) for f in os.listdir(pcb_bg_folder) if f.lower().endswith(('.png', '.jpg', '.bmp'))]
# 缓存结构改为：字典，key为图片路径，value为该图多块缓存区域列表
_bg_cache = {}

def get_center_near_crop(img, crop_h, crop_w, max_offset_ratio=0.5):
    """
    从图像中心附近截取一块区域，允许一定偏移，偏移范围为crop尺寸的max_offset_ratio比例。
    """
    h, w = img.shape[:2]
    center_y, center_x = h // 2, w // 2

    max_offset_y = int(crop_h * max_offset_ratio)
    max_offset_x = int(crop_w * max_offset_ratio)

    # 计算偏移范围，确保不越界
    y_min = max(center_y - max_offset_y, 0)
    y_max = min(center_y + max_offset_y, h - crop_h)
    x_min = max(center_x - max_offset_x, 0)
    x_max = min(center_x + max_offset_x, w - crop_w)

    if y_max < y_min:
        y_min, y_max = 0, h - crop_h
    if x_max < x_min:
        x_min, x_max = 0, w - crop_w

    y = random.randint(y_min, y_max)
    x = random.randint(x_min, x_max)

    return img[y:y+crop_h, x:x+crop_w]



def get_random_pcb_bg(target_h, target_w):
    """
    优化版：
    - 允许同一张图片缓存多块区域，增加多样性
    - 截取区域不必严格中心，只需靠近中心，允许一定偏移
    - 缓存区域尺寸为目标尺寸的3倍，保证更丰富的特征
    - 缓存数量有限，避免内存爆满
    """
    max_cache_regions_per_img = 6  # 每张图缓存最多3块区域
    max_cache_img_num =  10      # 最多缓存3张图

    while True:
        bg_path = random.choice(pcb_bg_paths)

        # 初始化缓存列表
        if bg_path not in _bg_cache:
            _bg_cache[bg_path] = []

        # 如果缓存区域不足，读取整图并缓存新区域
        if len(_bg_cache[bg_path]) < max_cache_regions_per_img:
            full_img = cv2.imread(bg_path, cv2.IMREAD_UNCHANGED)
            if full_img is None:
                print(f"Warning: failed to load PCB background image {bg_path}")
                continue
            h, w = full_img.shape[:2]

            crop_h = min(target_h * 6, h)
            crop_w = min(target_w * 6, w)

            new_crop = get_center_near_crop(full_img, crop_h, crop_w, max_offset_ratio=0.3)[:,:,:3]
            _bg_cache[bg_path].append(new_crop)

            # 控制缓存图片数量，超过限制则随机删除一张缓存及其所有区域
            if len(_bg_cache) > max_cache_img_num:
                remove_key = random.choice(list(_bg_cache.keys()))
                _bg_cache.pop(remove_key)

        # 从缓存中随机选一块区域
        cached_regions = _bg_cache[bg_path]
        if not cached_regions:
            continue
        center_crop = random.choice(cached_regions)

        ch, cw = center_crop.shape[:2]
        if ch < target_h or cw < target_w:
            # 缓存区域太小，跳过
            continue

        # 在缓存区域内随机截取目标大小块
        y = random.randint(0, ch - target_h)
        x = random.randint(0, cw - target_w)
        crop_bg = center_crop[y:y+target_h, x:x+target_w].copy()
        return crop_bg
"""
用于合成数据集
能够生成元件尺度的垂直bbox，而不是引脚尺度的
"""
def load_image(file_path):
    """Load a single image from the given path."""
    return cv2.imread(file_path, cv2.IMREAD_UNCHANGED)

def rotate_and_resize(image, mask, depth, vis,angle, scale):
    """Rotate and resize an image, its mask, and its depth map."""
    h, w = image.shape[:2]
    center = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D(center, angle, scale)

    # Compute the bounding box for the rotated image
    cos = np.abs(M[0, 0])
    sin = np.abs(M[0, 1])
    new_w = int((h * sin) + (w * cos))
    new_h = int((h * cos) + (w * sin))

    # Adjust the rotation matrix to account for the new dimensions
    M[0, 2] += (new_w / 2) - center[0]
    M[1, 2] += (new_h / 2) - center[1]

    # Perform the rotation and resizing
    rotated_image = cv2.warpAffine(image, M, (new_w, new_h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0, 0))
    rotated_mask = cv2.warpAffine(mask, M, (new_w, new_h), flags=cv2.INTER_NEAREST, borderMode=cv2.BORDER_CONSTANT, borderValue=0)
    rotated_depth = cv2.warpAffine(depth, M, (new_w, new_h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=0)
    rotated_vis = cv2.warpAffine(vis, M, (new_w, new_h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0, 0))

    return rotated_image, rotated_mask, rotated_depth,rotated_vis

def is_non_overlapping(mask, component_mask, x, y):
    """Check if a component can be placed at (x, y) without overlapping."""
    ch, cw = component_mask.shape[:2]
    if x + cw > mask.shape[1] or y + ch > mask.shape[0]:
        return False

    region = mask[y:y + ch, x:x + cw]
    region_non_zero = region > 0
    component_mask_non_zero = component_mask > 0
    overlap = np.sum(np.logical_and(region_non_zero, component_mask_non_zero))
    # overlap = np.sum(cv2.bitwise_and(region > 0, component_mask > 0))
    return overlap == 0

def place_component(board, board_depth, board_mask, board_cpimask,board_vis_mask,
                    component, component_depth, component_mask,cpmask, component_vis_mask,
                    x, y):
    """Place a component on the board and update depth, mask, and visualized mask."""
    ch, cw = component.shape[:2]
    foreground_indices = component_mask > 0

    

    for c in range(3):
        board[y:y + ch, x:x + cw, c][foreground_indices] = component[:, :, c][foreground_indices]

    board_depth[y:y + ch, x:x + cw][foreground_indices] = component_depth[foreground_indices]

    board_mask[y:y + ch, x:x + cw][foreground_indices] = np.maximum(
        board_mask[y:y + ch, x:x + cw][foreground_indices], component_mask[foreground_indices]
    )
    board_cpimask[y:y + ch, x:x + cw][foreground_indices] = np.maximum(
        board_cpimask[y:y + ch, x:x + cw][foreground_indices], cpmask[foreground_indices]
    )

    board_vis_mask[y:y + ch, x:x + cw][foreground_indices] = component_vis_mask[foreground_indices]

    
    # board_vis_mask[y:y + ch, x:x + cw][foreground_indices] = cv2.applyColorMap(
    #     component_mask[foreground_indices].astype(np.uint8) * (255 // 178), cv2.COLORMAP_JET
    # )[foreground_indices]


def cpextract(component_root,component_depth_root,output_folder, min_scale=0.5, max_scale=1.5):
    """Generate synthetic PCB dataset."""
    def get_bbox_from_mask(mask,bgwh,shift):
        """
        从掩码图像中获取掩码部分的边界框。

        :param mask: 输入的掩码图像，形状为 (h, w)
        :return: 边界框信息 (x, y, width, height) xywh
        """
        # 找到掩码中非零像素的坐标
        rows, cols = np.nonzero(mask)
        bw,bh=bgwh
        sx,sy=shift

        if len(rows) == 0 or len(cols) == 0:
            return (0, 0, 0, 0)

        # 计算边界框的左上角坐标和宽高
        x = np.min(cols)
        y = np.min(rows)
        width = np.max(cols) - x + 1
        height = np.max(rows) - y + 1

        return ((x+width/2+sx)/bw, (y+height/2+sy)/bh, width/bw, height/bh)
    
    def get_component_type(component_path):
        type_name=component_path[len(component_root):]
        type_name=os.path.dirname(type_name)
        pattern = r'JRS(\d+)[a-zA-Z].*'
        match = re.match(pattern, type_name)
        number = match.group(1)
        return int(number)
    
    # Load the empty board and depth
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    component_paths = []#len=1948
    # for root, dirs, files in os.walk(component_root):
    #     for file in files:
    #         if file.endswith(('.png', '.jpg', '.jpeg')) and 'label' not in root:
    #             component_paths.append(os.path.join(root, file))
    #注意component_paths的id不是image的class_id，是乱序的
    # sub_dirs = sorted([sub_dir for sub_dir in os.listdir(component_root) if os.path.isdir(os.path.join(component_root, sub_dir))])
    for sub_dir in os.listdir(component_root):
    # for sub_dir in sub_dirs:
        sub_dir_path = os.path.join(component_root, sub_dir)
        # 检查是否为文件夹
        if os.path.isdir(sub_dir_path):
            image_paths = []
            # 遍历子文件夹内的所有文件
            for file in os.listdir(sub_dir_path):
                file_path = os.path.join(sub_dir_path, file)
                # 检查文件是否为图像文件（这里以常见的图像格式为例）
                if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')) and 'label' not in sub_dir_path and ('SolderLight' in file):
                    error_flag=False
                    for error_path in error_guest:
                        if error_path in file_path:
                        
                            error_flag=True
                            break
                    if error_flag:
                        continue
                    image_paths.append(file_path)
            # 将当前子文件夹的图像路径列表添加到二维列表中
            if image_paths:
                component_paths.append(image_paths)

    process_images(component_paths, output_folder,component_depth_root)


def crop_and_resize_image(image, mask,depth):
    """
    根据分割掩码裁剪图像并调整大小到 100x100
    :param image: 输入的图像
    :param mask: 对应的分割掩码
    :return: 裁剪并调整大小后的图像和掩码
    """
    # 找到掩码中非零区域的边界
    y_nonzero, x_nonzero = np.nonzero(mask)
    min_y, max_y = np.min(y_nonzero), np.max(y_nonzero)
    min_x, max_x = np.min(x_nonzero), np.max(x_nonzero)

    # 裁剪图像和掩码
    cropped_image = image[min_y:max_y + 1, min_x:max_x + 1]
    cropped_mask = mask[min_y:max_y + 1, min_x:max_x + 1]
    cropped_depth=depth[min_y:max_y + 1, min_x:max_x + 1]

    # 创建一个全黑的背景图像，尺寸与cropped_image相同
    # black_background = np.zeros_like(cropped_image)
    ih,iw=cropped_image.shape[:2]
    pcb_background=get_random_pcb_bg(ih,iw)
    # 将掩码应用到图像上
    result = cv2.bitwise_and(cropped_image, cropped_image, mask=cropped_mask)
    # 将结果复制到黑色背景上
    pcb_background[np.where((result != [0, 0, 0]).all(axis=2))] = result[
        np.where((result != [0, 0, 0]).all(axis=2))
    ]
    cropped_image=pcb_background
    # 调整大小到 100x100
    resized_image = cv2.resize(cropped_image, (CP_IMG_SIZE, CP_IMG_SIZE),interpolation=cv2.INTER_NEAREST)
    resized_mask = cv2.resize(cropped_mask, (CP_IMG_SIZE, CP_IMG_SIZE),interpolation=cv2.INTER_NEAREST)
    resized_depth=cv2.resize(cropped_depth,(CP_IMG_SIZE,CP_IMG_SIZE),interpolation=cv2.INTER_NEAREST)

    # 创建一个布尔数组，标记出不等于 0、1、2 的元素
    # non_target_mask = (resized_mask != 0) & (resized_mask != 1) & (resized_mask != 2)
    # # 将标记为 True 的元素赋值为 3
    # resized_mask[non_target_mask] = 3

    del image,mask,depth,cropped_depth,cropped_image,cropped_mask
    return resized_image, resized_mask,resized_depth

def process_images(image_paths, base_output_path, component_depth_root):
    """
    处理图像、掩码和深度图，裁剪、调整大小并保存到指定路径
    :param image_paths: 二维列表，包含所有元件图片的路径
    :param base_output_path: 输出的基础路径
    :param component_depth_root: 深度图的根路径
    """
    # 创建 images、masks 和 depths 文件夹
    images_output_path = os.path.join(base_output_path, 'images')
    masks_output_path = os.path.join(base_output_path,'masks')
    depths_output_path = os.path.join(base_output_path, 'depths')
    os.makedirs(images_output_path, exist_ok=True)
    os.makedirs(masks_output_path, exist_ok=True)
    os.makedirs(depths_output_path, exist_ok=True)

    for sub_list in image_paths:
        for image_path in sub_list:
            # 获取对应的分割掩码路径
            # mask_path = (image_path).replace(os.path.basename(os.path.dirname(image_path)), 'label').replace('jpg','png')
            mask_path=os.path.join(os.path.join(os.path.dirname(image_path),'label'),os.path.basename(image_path).replace('jpg','png'))
            # 获取对应的深度图路径
            depth_path = os.path.join(os.path.join(component_depth_root,os.path.basename(os.path.dirname(image_path))), os.path.basename(image_path).replace('_SolderLight.jpg','.png'))

            # 读取图像、掩码和深度图
            image = cv2.imread(image_path)
            mask = cv2.imread(mask_path.replace('SolderLight','TrueLight'), cv2.IMREAD_GRAYSCALE)
            depth = cv2.imread(depth_path, cv2.IMREAD_ANYDEPTH)

            if depth is None:
                print('here')
            # 裁剪并调整大小
            resized_image, resized_mask,resized_depth = crop_and_resize_image(image, mask,depth)
            # _, resized_depth = crop_and_resize_image(, mask)

            # 获取文件名
            file_name = os.path.basename(image_path)

            # 保存处理后的图像、掩码和深度图
            cv2.imwrite(os.path.join(images_output_path, file_name).replace('jpg','png'), resized_image)
            cv2.imwrite(os.path.join(masks_output_path, file_name).replace('jpg','png'), resized_mask)
            cv2.imwrite(os.path.join(depths_output_path, file_name).replace('jpg','png'), resized_depth)

error_guest=[
            "JRS183R0603/_a_SaveData__20241013161630397_Save Data [14837, 35647]",
            "JRS156R0603/_a_SaveData__20241013161630396_Save Data [22823, 35559]",
            "JRS157R0603/_a_SaveData__20241013161630397_Save Data [18766, 35613]",
            "JRS116PGLQFP/_a_SaveData__20241013155754673_SaveData [23447, 102811]",
            "JRS104TSOP50",
            "JRS116PGLQFP",
            "JRS82R1812/_a_SaveData__20241013155402990_Save Data [109592, 18847]",
        ]
component_root="/home/lhx/pcbsyn/data_zoo/PCBA/20241209_PCBA样板标图数据含真彩-华科用/"
output_folder="data_zoo/pcb_truecolor_subcpjrsbg"
component_depth_root="/home/lhx/pcbsyn/pcbad"

#32*8=256
CP_IMG_SIZE=640
cpextract(
    component_root=component_root,
    component_depth_root=component_depth_root,
    output_folder=output_folder,
    min_scale=1.1,
    max_scale=0.9
)


































